package com.fx.zmlzml.filter;

import org.jsoup.Jsoup;
import org.jsoup.safety.Safelist;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Pattern;

/**
 * XSS请求包装器
 * 用于清理HTTP请求中的各种参数，防止XSS攻击
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    // 使用Jsoup库提供的安全列表进行HTML清理
    private static final Safelist SAFE_LIST = Safelist.basicWithImages()
            .addAttributes("a", "href")
            .addAttributes("img", "src", "alt")
            .addTags("p", "br", "strong", "em", "u", "ul", "ol", "li");

    // 用于清理JavaScript代码的正则表达式
    private static final Pattern SCRIPT_PATTERN = Pattern.compile(
            "<script\\b[^>]*>(.*?)</script>", 
            Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL
    );

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        return value != null ? cleanXss(value) : null;
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values != null) {
            for (int i = 0; i < values.length; i++) {
                values[i] = cleanXss(values[i]);
            }
        }
        return values;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        Map<String, String[]> cleanMap = new HashMap<>();

        for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
            String[] values = entry.getValue();
            String[] cleanValues = new String[values.length];
            
            for (int i = 0; i < values.length; i++) {
                cleanValues[i] = cleanXss(values[i]);
            }
            
            cleanMap.put(entry.getKey(), cleanValues);
        }
        
        return cleanMap;
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        return value != null ? cleanXss(value) : null;
    }

    /**
     * 清理字符串中的XSS内容
     * @param value 需要清理的字符串
     * @return 清理后的安全字符串
     */
    private String cleanXss(String value) {
        if (value == null) {
            return null;
        }
        
        // 第一步：使用正则表达式移除<script>标签
        String cleanValue = SCRIPT_PATTERN.matcher(value).replaceAll("");
        
        // 第二步：使用Jsoup库清理HTML标签和属性
        cleanValue = Jsoup.clean(cleanValue, SAFE_LIST);
        
        // 第三步：对特殊字符进行编码
        cleanValue = escapeHtml(cleanValue);
        
        return cleanValue;
    }

    /**
     * 对HTML特殊字符进行编码
     * @param input 需要编码的字符串
     * @return 编码后的字符串
     */
    private String escapeHtml(String input) {
        if (input == null) {
            return null;
        }
        
        return input
                .replace("&", "&amp;")
                .replace("<", "&lt;")
                .replace(">", "&gt;")
                .replace("\"", "&quot;")
                .replace("'", "&#x27;")
                .replace("/", "&#x2F;");
    }
}